#include <iostream>
#include <opencv2/opencv.hpp>
#include "bsnn_model_load.h"
#include "drmshow.h"
#include "image_process.h"
#include <string.h>


using namespace std;
std::string input_path = "/workspace/input/fad_quick_start_video_02.avi";
// std::string input_path = "/workspace/input/fad_video.jpg";
std::string model_path = "/workspace/models/yolov5/";



int main(int argc, char* argv[])
{
    cv::Mat img, input_image;
    cv::VideoCapture capture(input_path);
    BSNN_MODEL bsnn_model(model_path);

    while (true)
    {
    // step 1 图片预处理
    Timer time;
    
    capture >> img;
    time.reset();
    preprocess(img, input_image);
    
    size_t input_img_len = 3 * IMG_HEIGHT * IMG_WIDTH;

    // 更改图片在内存中的存储方式为WHD
    if (!input_image.isContinuous())
        printf("-> input image is not continuous in memory...");
    uchar input_buf[input_img_len] = {0};
    for (int c = 0; c < 3; c++)
    {
        for (int i = 0; i < IMG_HEIGHT; i++)
        {
            for (int j = 0; j < IMG_WIDTH; j++)
            {
                input_buf[c * IMG_HEIGHT * IMG_WIDTH + i * IMG_WIDTH + j] = input_image.data[3 * (i * IMG_WIDTH + j) + c];
            }
        }
    }
    cout << "-> preprocess time : " << time.elapsed() << endl;

    // step 2 模型推理
    time.reset();
    bsnn_model.Run(input_buf, input_img_len);
    auto bsnn_output = bsnn_model.GetModelOutput();
    cout << "-> bsnn model inference time : " << time.elapsed() << endl;

    // step 3 后处理
    time.reset();
    std::vector<ObjInfo> result;
    process_output(bsnn_output.get(), result);
    // cout << "-> The number of objects detected: " << result.size() << endl;
    cout << "-> post process time : " << time.elapsed() << endl;

    draw_bboxes(img, result);
    bsnn_model.ReleaseOutputBuffer();


    cv::imshow("img", img);
    cv::waitKey(1);
    }


    return 0;
}